# -*- coding: UTF-8 -*-
"""
@Date    ：2025/9/11 17:26
@Author  ：Liu Yuezhao
@Project ：bert
@File    ：time2vec.py
@IDE     ：PyCharm
"""
import torch
import torch.nn as nn

class Time2Vec(nn.Module):
    """
    Time2Vec for a single time feature
    z_k = sin(w_k * t + p_k) for k=0, ..., K-1
    z_K = w_K * t + p_K
    """
    def __init__(self, embed_dim, activation='sin'):
        super().__init__()
        self.embed_dim = embed_dim
        self.activation = activation
        self.w = nn.Parameter(torch.randn(1, 1, embed_dim - 1))
        self.b = nn.Parameter(torch.randn(1, 1, embed_dim - 1))
        self.w_a = nn.Parameter(torch.ones(1, 1, 1))
        self.b_a = nn.Parameter(torch.zeros(1, 1, 1))

    def forward(self, x):
        x = x.unsqueeze(-1)  # (B, L, 1)
        v1 = self.w * x + self.b  # (B, L, D-1)
        if self.activation == "sin":
            v1 = torch.sin(v1)
        elif self.activation == "cos":
            v1 = torch.cos(v1)
        else:
            raise ValueError("activation must be 'sin' or 'cos'")
        v2 = self.w_a * x + self.b_a  # (B, L, 1)
        return torch.cat([v2, v1], dim=-1)  # (B, L, D)


class TripleTime2Vec(nn.Module):
    """
    分别对三种时间信号进行编码
    """
    def __init__(self, embed_dim=768, activation='cos'):
        super().__init__()
        self.embed_dim = embed_dim
        self.activation = activation

        # 为每个时间信号独立编码
        self.interval_encoder = Time2Vec(embed_dim, activation)
        self.same_event_encoder = Time2Vec(embed_dim, activation)
        self.age_encoder = Time2Vec(embed_dim, activation)

    def forward(self, interval, same_event_interval, age):
        """
        Args:
            interval: (B, L), 相邻事件时间差
            same_event_interval: (B, L), 同类事件时间差
            age: (B, L), 距离序列起点时间
        Returns:
            interval_embed: (B, L, D)
            same_event_embed: (B, L, D)
            age_embed: (B, L, D)
        """
        return (
            self.interval_encoder(interval),
            self.same_event_encoder(same_event_interval),
            self.age_encoder(age)
        )